#pragma once

#include "concepts.h"
#include "previous.h"
#include <coroutine>
#include <exception>
#include <iostream>
#include <optional>

namespace co_async {
template <class T> struct GeneratorPromise {
  auto initial_suspend() noexcept { return std::suspend_always{}; }
  auto final_suspend() noexcept { return PreviousAwaiter{m_previous}; }

  auto get_return_object() {
    return std::coroutine_handle<GeneratorPromise>::from_promise(*this);
  }

  void unhandled_exception() noexcept {
    std::cout << "unhandled_exception" << std::endl;
    m_exception = std::current_exception();
   m_final = true;
  }

  auto yield_value(T &&val) {
    m_result.put(std::move(val));
    return PreviousAwaiter{m_previous};
  }

  auto yield_value(const T &val) {
    m_result.put(val);
    return PreviousAwaiter{m_previous};
  }

  bool final() {
    if (m_final) {
      if (m_exception) [[unlikely]] {
        std::rethrow_exception(m_exception);
      }
    }
    return m_final;
  }

  void return_void() { m_final = true; }

  auto result() { return m_result.move(); }

  std::coroutine_handle<> m_previous;
  std::exception_ptr m_exception;
  UnInit<T> m_result;
  bool m_final;

  GeneratorPromise &operator=(GeneratorPromise &&) = delete;
};

template <class T> struct GeneratorPromise<T &> {
  auto initial_suspend() noexcept { return std::suspend_always{}; }
  auto final_suspend() noexcept { return PreviousAwaiter{m_previous}; }

  auto get_return_object() {
    return std::coroutine_handle<GeneratorPromise>::from_promise(*this);
  }

  void unhandled_exception() noexcept {
    std::cout << "unhandled_exception" << std::endl;
    m_exception = std::current_exception();
    m_result = nullptr;
  }

  auto yield_value(T &val) {
    m_result = &val;
    return PreviousAwaiter{m_previous};
  }

  bool final() {
    if (!m_result) {
      if (m_exception) [[unlikely]] {
        std::rethrow_exception(m_exception);
      }
      return true;
    }
    return false;
  }

  void return_void() { m_result = nullptr; }

  T &result() { return *m_result; }

  std::coroutine_handle<> m_previous;
  std::exception_ptr m_exception;
  T *m_result;

  GeneratorPromise &operator=(GeneratorPromise &&) = delete;
};

template <class T, class P = GeneratorPromise<T>> struct Generator {
  using promise_type = P;
  Generator(Generator &&) = delete;
  Generator(std::coroutine_handle<promise_type> coroutine = nullptr) noexcept
        : m_handle(coroutine) {}

  ~Generator() { m_handle.destroy(); }

  struct Awaiter {
    bool await_ready() const noexcept { return false; }

    std::coroutine_handle<promise_type>
    await_suspend(std::coroutine_handle<> coroutine) const noexcept {
      m_handle.promise().m_previous = coroutine;
      return m_handle;
    }

    std::optional<T> await_resume() const {
      if (m_handle.promise().final())
        return std::nullopt;
      return m_handle.promise().result();
    }

    std::coroutine_handle<promise_type> m_handle;
  };

  auto operator co_await() const noexcept { return Awaiter(m_handle); }

  operator std::coroutine_handle<promise_type>() const noexcept {
    return m_handle;
  }

private:
  std::coroutine_handle<promise_type> m_handle;
};

} // namespace co_async